import math
import torch
import numpy as np

from torch.optim.lr_scheduler import _LRScheduler

class ExpUpCosDown(_LRScheduler):
    def __init__(
            self, 
            optimizer, 
            lrs, 
            warm_steps, 
            exponent=-5.0, 
            tot_steps=None, 
            min_lr=1e-6, 
            mode="ExpUpCosDown",
            joint_warmup=0.0
        ):
        self.opt = optimizer
        self.rampup_len = warm_steps
        self.max_lrs = lrs
        self.exponent = exponent
        self.tot_steps = tot_steps
        self.min_lr = min_lr
        self.mode = mode
        self.joint_warmup = joint_warmup
        assert mode in ["ExpUpCosDown", "ExpUp"], "Unknown mode: {}".format(mode)
        super(ExpUpCosDown, self).__init__(optimizer)

    def _get_scaling_factor(self):
        if self.rampup_len == 0:
            return 1.0
        elif self._step_count < self.rampup_len:
            current = np.clip(self._step_count, 0.0, self.rampup_len)
            phase = 1.0 - current / self.rampup_len
            return float(np.exp(self.exponent * phase * phase))
        elif self.mode == "ExpUpCosDown": 
            current = np.clip(self._step_count, self.rampup_len, self.tot_steps)
            phase = (current - self.rampup_len) / (self.tot_steps - self.rampup_len) * math.pi
            return (math.cos(phase) + 1) / 2
        elif self.mode == "ExpUp":
            return 1.0
    
    # cosine down for SelfSL joint weight
    def _get_interpolate_factor(self):
        current = np.clip(self._step_count, 0, self.tot_steps)
        if current < self.joint_warmup:
            phase = (current) / (self.joint_warmup) * math.pi
        else:
            phase = math.pi
        return (math.cos(phase) + 1) / 2

    def get_lr(self):
        if self._step_count < self.rampup_len:
            return [lr * self._get_scaling_factor() for lr in self.max_lrs]
        elif self.mode == "ExpUpCosDown":
            return [self.min_lr + (lr - self.min_lr) * self._get_scaling_factor() for lr in self.max_lrs]
        elif self.mode == "ExpUp":
            return self.max_lrs